import tkinter as tk
from tkinter import ttk, messagebox
import paramiko
import threading
import socket
import time
import sys
import json
import os

class SSHPortForwarder:
    def __init__(self, root):
        self.root = root
        self.root.title("SSH Port Forwarder")
        self.root.geometry("600x550")
        
        # 存储端口转发连接信息
        self.forwarding_connections = {}
        self.is_running = False
        
        # 创建界面元素
        self.create_widgets()
        
        # 加载配置文件（如果存在）
        self.load_config()
        
    def create_widgets(self):
        # 主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 配置网格权重
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        main_frame.columnconfigure(1, weight=1)
        
        # 远程主机信息区域
        host_frame = ttk.LabelFrame(main_frame, text="远程主机配置", padding="10")
        host_frame.grid(row=0, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=(0, 10))
        host_frame.columnconfigure(1, weight=1)
        
        # IP地址
        ttk.Label(host_frame, text="IP地址:").grid(row=0, column=0, sticky=tk.W, padx=(0, 5), pady=2)
        self.ip_entry = ttk.Entry(host_frame, width=30)
        self.ip_entry.grid(row=0, column=1, sticky=(tk.W, tk.E), padx=(0, 10), pady=2)
        self.ip_entry.insert(0, "192.168.1.1")
        
        # 端口
        ttk.Label(host_frame, text="端口:").grid(row=0, column=2, sticky=tk.W, padx=(0, 5), pady=2)
        self.port_entry = ttk.Entry(host_frame, width=10)
        self.port_entry.grid(row=0, column=3, sticky=tk.W, padx=(0, 10), pady=2)
        self.port_entry.insert(0, "22")
        
        # 用户名
        ttk.Label(host_frame, text="用户名:").grid(row=1, column=0, sticky=tk.W, padx=(0, 5), pady=2)
        self.username_entry = ttk.Entry(host_frame, width=30)
        self.username_entry.grid(row=1, column=1, sticky=(tk.W, tk.E), padx=(0, 10), pady=2)
        
        # 密码
        ttk.Label(host_frame, text="密码:").grid(row=1, column=2, sticky=tk.W, padx=(0, 5), pady=2)
        self.password_entry = ttk.Entry(host_frame, width=30, show="*")
        self.password_entry.grid(row=1, column=3, sticky=(tk.W, tk.E), padx=(0, 10), pady=2)
        
        # 端口映射区域
        mapping_frame = ttk.LabelFrame(main_frame, text="端口映射", padding="10")
        mapping_frame.grid(row=1, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N, tk.S), pady=(0, 10))
        mapping_frame.columnconfigure(0, weight=1)
        mapping_frame.rowconfigure(0, weight=1)
        
        # 监听端口
        ttk.Label(mapping_frame, text="监听端口:").grid(row=0, column=0, sticky=tk.W, padx=(0, 5), pady=2)
        self.local_port_entry = ttk.Entry(mapping_frame, width=10)
        self.local_port_entry.grid(row=0, column=1, sticky=tk.W, padx=(0, 10), pady=2)
        
        # 目标主机
        ttk.Label(mapping_frame, text="目标主机:").grid(row=0, column=2, sticky=tk.W, padx=(0, 5), pady=2)
        self.remote_host_entry = ttk.Entry(mapping_frame, width=20)
        self.remote_host_entry.grid(row=0, column=3, sticky=tk.W, padx=(0, 10), pady=2)
        
        # 目标端口
        ttk.Label(mapping_frame, text="目标端口:").grid(row=0, column=4, sticky=tk.W, padx=(0, 5), pady=2)
        self.remote_port_entry = ttk.Entry(mapping_frame, width=10)
        self.remote_port_entry.grid(row=0, column=5, sticky=tk.W, padx=(0, 10), pady=2)
        
        # 按钮区域
        button_frame = ttk.Frame(main_frame)
        button_frame.grid(row=2, column=0, columnspan=2, pady=(0, 10))
        
        self.start_button = ttk.Button(button_frame, text="开始运行", command=self.start_forwarding)
        self.start_button.pack(side=tk.LEFT, padx=(0, 5))
        
        self.stop_button = ttk.Button(button_frame, text="结束运行", command=self.stop_forwarding, state=tk.DISABLED)
        self.stop_button.pack(side=tk.LEFT, padx=(0, 5))
        
        # 添加保存配置按钮
        save_button = ttk.Button(button_frame, text="保存配置", command=self.save_config)
        save_button.pack(side=tk.LEFT, padx=(0, 5))
        
        # 端口映射列表
        list_frame = ttk.LabelFrame(main_frame, text="端口映射列表", padding="10")
        list_frame.grid(row=3, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N, tk.S), pady=(0, 10))
        list_frame.columnconfigure(0, weight=1)
        list_frame.rowconfigure(0, weight=1)
        
        # 创建Treeview显示端口映射列表
        columns = ("监听端口", "目标主机", "目标端口", "状态")
        self.mapping_listbox = ttk.Treeview(list_frame, columns=columns, show="headings", height=8)
        
        for col in columns:
            self.mapping_listbox.heading(col, text=col)
            self.mapping_listbox.column(col, width=100)
            
        # 添加滚动条
        scrollbar_y = ttk.Scrollbar(list_frame, orient=tk.VERTICAL, command=self.mapping_listbox.yview)
        scrollbar_x = ttk.Scrollbar(list_frame, orient=tk.HORIZONTAL, command=self.mapping_listbox.xview)
        
        self.mapping_listbox.configure(yscrollcommand=scrollbar_y.set, xscrollcommand=scrollbar_x.set)
        
        # 布局
        self.mapping_listbox.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        scrollbar_y.grid(row=0, column=1, sticky=(tk.N, tk.S))
        scrollbar_x.grid(row=1, column=0, sticky=(tk.W, tk.E))
        
        # 添加按钮
        add_button = ttk.Button(main_frame, text="添加映射", command=self.add_mapping)
        add_button.grid(row=4, column=0, pady=(5, 0), sticky=tk.W)
        
        remove_button = ttk.Button(main_frame, text="删除选中", command=self.remove_mapping)
        remove_button.grid(row=4, column=1, pady=(5, 0), sticky=tk.W)
        
    def add_mapping(self):
        """添加端口映射"""
        try:
            local_port = int(self.local_port_entry.get())
            remote_host = self.remote_host_entry.get()
            remote_port = int(self.remote_port_entry.get())
            
            if not remote_host or not remote_port:
                messagebox.showerror("错误", "请填写完整的目标主机和端口信息")
                return
                
            # 检查是否已存在相同的映射
            for item in self.mapping_listbox.get_children():
                values = self.mapping_listbox.item(item)['values']
                if (values[0] == local_port and 
                    values[1] == remote_host and 
                    values[2] == remote_port):
                    messagebox.showwarning("警告", "该端口映射已存在")
                    return
            
            # 添加到列表
            self.mapping_listbox.insert("", tk.END, values=(local_port, remote_host, remote_port, "未运行"))
            
            # 清空输入框
            self.local_port_entry.delete(0, tk.END)
            self.remote_host_entry.delete(0, tk.END)
            self.remote_port_entry.delete(0, tk.END)
            
        except ValueError:
            messagebox.showerror("错误", "端口必须是数字")
    
    def remove_mapping(self):
        """删除选中的映射"""
        selected = self.mapping_listbox.selection()
        if not selected:
            messagebox.showwarning("警告", "请先选择要删除的映射")
            return
            
        for item in selected:
            self.mapping_listbox.delete(item)
            
    def start_forwarding(self):
        """开始端口转发"""
        # 保存当前配置先
        self.save_config()
        
        # 获取远程主机信息
        hostname = self.ip_entry.get()
        port = int(self.port_entry.get())
        username = self.username_entry.get()
        password = self.password_entry.get()
        
        if not all([hostname, username, password]):
            messagebox.showerror("错误", "请填写完整的远程主机信息")
            return
            
        # 检查是否有映射
        if len(self.mapping_listbox.get_children()) == 0:
            messagebox.showwarning("警告", "请至少添加一个端口映射")
            return
            
        try:
            # 启动SSH连接和转发线程
            self.ssh_thread = threading.Thread(target=self.run_forwarding, 
                                              args=(hostname, port, username, password))
            self.ssh_thread.daemon = True
            self.ssh_thread.start()
            
            # 更新按钮状态
            self.start_button.config(state=tk.DISABLED)
            self.stop_button.config(state=tk.NORMAL)
            self.is_running = True
            
        except Exception as e:
            messagebox.showerror("错误", f"启动失败: {str(e)}")
    
    def stop_forwarding(self):
        """停止端口转发"""
        # 停止所有连接
        for connection in self.forwarding_connections.values():
            try:
                if hasattr(connection, 'close'):
                    connection.close()
            except Exception as e:
                print(f"关闭连接时出错: {e}")
                
        # 清空连接列表
        self.forwarding_connections.clear()
        
        # 更新按钮状态
        self.start_button.config(state=tk.NORMAL)
        self.stop_button.config(state=tk.DISABLED)
        self.is_running = False
        
        # 更新列表显示为停止状态
        for item in self.mapping_listbox.get_children():
            values = list(self.mapping_listbox.item(item)['values'])
            if values[3] == "运行中":
                values[3] = "已停止"
                self.mapping_listbox.item(item, values=values)
                
    def run_forwarding(self, hostname, port, username, password):
        """实际执行端口转发的线程函数"""
        try:
            # 创建SSH客户端
            ssh_client = paramiko.SSHClient()
            ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # 连接远程主机
            ssh_client.connect(hostname, port=port, username=username, password=password)
            
            # 获取所有映射并建立转发连接
            mappings = []
            for item in self.mapping_listbox.get_children():
                values = self.mapping_listbox.item(item)['values']
                mappings.append(values)
                
            # 更新列表显示为运行中状态
            for i, (local_port, remote_host, remote_port, status) in enumerate(mappings):
                if not self.is_running:
                    break
                    
                try:
                    # 创建端口转发线程
                    forward_thread = threading.Thread(
                        target=self.create_forwarding_connection,
                        args=(ssh_client, local_port, remote_host, remote_port)
                    )
                    forward_thread.daemon = True
                    forward_thread.start()
                    
                    # 更新列表显示为运行中状态
                    self.root.after(0, lambda i=i: self.update_mapping_status(i, "运行中"))
                    
                except Exception as e:
                    print(f"建立转发连接失败 {local_port} -> {remote_host}:{remote_port}: {e}")
                    # 更新列表显示为错误状态
                    self.root.after(0, lambda i=i: self.update_mapping_status(i, "错误"))
                    
            # 等待线程结束（实际应用中可能需要更复杂的逻辑）
            while self.is_running:
                time.sleep(1)
                
        except Exception as e:
            messagebox.showerror("错误", f"SSH连接失败: {str(e)}")
            
    def create_forwarding_connection(self, ssh_client, local_port, remote_host, remote_port):
        """创建端口转发连接"""
        try:
            # 创建本地socket监听
            local_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            local_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            local_socket.bind(('localhost', local_port))
            local_socket.listen(5)
            
            # 将连接信息保存到字典中，用于后续关闭
            connection_key = f"{local_port}_{remote_host}_{remote_port}"
            self.forwarding_connections[connection_key] = {
                'socket': local_socket,
                'ssh_client': ssh_client
            }
            
            print(f"监听端口 {local_port} 转发到 {remote_host}:{remote_port}")
            
            # 持续接受连接并转发数据
            while self.is_running:
                try:
                    client_socket, addr = local_socket.accept()
                    print(f"新连接来自: {addr}")
                    
                    # 创建远程连接线程
                    threading.Thread(
                        target=self.handle_connection,
                        args=(ssh_client, client_socket, remote_host, remote_port),
                        daemon=True
                    ).start()
                except Exception as e:
                    if self.is_running:
                        print(f"接受连接时出错: {e}")
                        
        except Exception as e:
            print(f"创建转发连接失败: {e}")
            
    def handle_connection(self, ssh_client, client_socket, remote_host, remote_port):
        """处理客户端连接"""
        try:
            # 创建到远程主机的SSH隧道
            transport = ssh_client.get_transport()
            dest_addr = (remote_host, remote_port)
            local_addr = ('localhost', 0)  # 使用系统分配的本地端口
            
            # 打开通道进行转发
            channel = transport.open_channel("direct-tcpip", dest_addr, local_addr)
            
            if not channel:
                print(f"无法打开到 {remote_host}:{remote_port} 的通道")
                client_socket.close()
                return
                
            # 双向数据传输
            def forward_data(src, dst):
                try:
                    while self.is_running:
                        data = src.recv(1024)
                        if not data:
                            break
                        dst.send(data)
                except Exception as e:
                    pass
                finally:
                    try:
                        src.close()
                        dst.close()
                    except:
                        pass
                        
            # 启动双向转发线程
            threading.Thread(target=forward_data, args=(client_socket, channel), daemon=True).start()
            threading.Thread(target=forward_data, args=(channel, client_socket), daemon=True).start()
            
        except Exception as e:
            print(f"处理连接时出错: {e}")
            try:
                client_socket.close()
            except:
                pass
                
    def update_mapping_status(self, index, status):
        """更新映射状态显示"""
        try:
            item = self.mapping_listbox.get_children()[index]
            values = list(self.mapping_listbox.item(item)['values'])
            values[3] = status
            self.mapping_listbox.item(item, values=values)
        except IndexError:
            pass  # 索引超出范围，忽略
    
    def save_config(self):
        """保存配置到config.json文件"""
        config = {
            "host": {
                "ip": self.ip_entry.get(),
                "port": self.port_entry.get(),
                "username": self.username_entry.get(),
                # "password": self.password_entry.get()
            },
            "mappings": []
        }
        
        # 收集所有端口映射
        for item in self.mapping_listbox.get_children():
            values = self.mapping_listbox.item(item)['values']
            config["mappings"].append({
                "local_port": values[0],
                "remote_host": values[1],
                "remote_port": values[2]
            })
        
        try:
            with open("config.json", "w", encoding="utf-8") as f:
                json.dump(config, f, ensure_ascii=False, indent=4)
            print("配置已保存到 config.json")
        except Exception as e:
            messagebox.showerror("错误", f"保存配置失败: {str(e)}")
    
    def load_config(self):
        """从config.json文件加载配置"""
        if not os.path.exists("config.json"):
            return
            
        try:
            with open("config.json", "r", encoding="utf-8") as f:
                config = json.load(f)
            
            # 加载主机信息
            if "host" in config:
                host_config = config["host"]
                self.ip_entry.delete(0, tk.END)
                self.ip_entry.insert(0, host_config.get("ip", "192.168.1.1"))
                
                self.port_entry.delete(0, tk.END)
                self.port_entry.insert(0, host_config.get("port", "22"))
                
                self.username_entry.delete(0, tk.END)
                self.username_entry.insert(0, host_config.get("username", ""))
                
                # 注意：密码不从文件加载，保持安全
                # 为了用户体验，可以考虑将密码字段置空
            
            # 加载端口映射
            if "mappings" in config:
                for mapping in config["mappings"]:
                    local_port = mapping.get("local_port")
                    remote_host = mapping.get("remote_host")
                    remote_port = mapping.get("remote_port")
                    
                    # 检查是否已存在相同的映射
                    exists = False
                    for item in self.mapping_listbox.get_children():
                        values = self.mapping_listbox.item(item)['values']
                        if (values[0] == local_port and 
                            values[1] == remote_host and 
                            values[2] == remote_port):
                            exists = True
                            break
                    
                    if not exists:
                        # 添加到列表
                        self.mapping_listbox.insert("", tk.END, values=(local_port, remote_host, remote_port, "未运行"))
            
            print("配置已从 config.json 加载")
        except Exception as e:
            messagebox.showerror("错误", f"加载配置失败: {str(e)}")

def main():
    root = tk.Tk()
    app = SSHPortForwarder(root)
    
    # 设置窗口关闭事件处理
    def on_closing():
        if messagebox.askokcancel("退出", "确定要退出程序吗？"):
            # 保存配置时停止所有连接
            app.stop_forwarding()  
            # 保存当前配置
            app.save_config()
            root.destroy()
            
    root.protocol("WM_DELETE_WINDOW", on_closing)
    
    try:
        root.mainloop()
    except KeyboardInterrupt:
        print("\n程序被中断")
        sys.exit(0)

if __name__ == "__main__":
    main()